package rpc_service

import (
	"context"
	"github.com/chicken-team-outside/chicken_transmission/api_messages"
	"github.com/chicken-team-outside/chicken_transmission/proxy_server/middleware"
	"github.com/chicken-team-outside/chicken_transmission/proxy_server/model"
	"github.com/chicken-team-outside/chicken_transmission/proxy_server/setting"
	"github.com/chicken-team-outside/chicken_transmission/proxy_server/utils"
	"github.com/chicken-team-outside/chicken_transmission/socket"
	"github.com/gin-gonic/gin"
	"github.com/golang/protobuf/proto"
	"github.com/gorilla/websocket"
	"go.uber.org/zap"
	"io"
	"sync"
	"time"
)

var clients map[uint64]*socket.WebSocketConnection
var upgrader *websocket.Upgrader
var clientMapLock sync.RWMutex

func init() {
	clients = make(map[uint64]*socket.WebSocketConnection)
}

func Init(group *gin.RouterGroup) {
	upgrader = &websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
	group.GET("/msg", handleConnectRequest)
}

func sendMsgToClient(ctx context.Context, messageId, clientId uint64, msgType api_messages.MessageType, msg proto.Message) error {
	client := GetClient(clientId)
	if client != nil {
		header := &api_messages.MessageHeader{
			Reply:     messageId,
			Type:      msgType,
			ExpiredAt: time.Now().Unix() + int64(setting.RpcSendTimeout.Seconds()),
		}
		var cancel context.CancelFunc
		ackCtx, cancel := context.WithTimeout(ctx, setting.RpcSendTimeout)
		defer cancel()
		return client.SendProtoMessage(ackCtx, header, msg)
	}
	return nil
}

func broadcastMsgToAllClientByUser(excludeClient, usrId uint64, ctx context.Context, msgType api_messages.MessageType, msg proto.Message) {
	for clientId := range model.GetRelations(usrId) {
		if clientId == excludeClient {
			continue
		}
		err := sendMsgToClient(ctx, 0, clientId, msgType, msg)
		if err != nil {
			zap.L().Debug("broadcast msg to client failed", zap.Error(err), zap.Uint64("clientId", clientId), zap.Int32("type", int32(msgType)))
		}
	}
}

func handleConnectRequest(ctx *gin.Context) {
	info := middleware.GetTransportInfo(ctx)
	if info == nil {
		ctx.JSON(403, &socket.ConnectionFailedResponse{
			Code: 403,
			Msg:  "invalid token",
		})
	} else {
		conn, err := upgrader.Upgrade(ctx.Writer, ctx.Request, nil)
		if err == nil {
			client := socket.NewWebSocket(conn, context.Background(), func() {
				clientMapLock.Lock()
				delete(clients, info.Id)
				clientMapLock.Unlock()
				zap.L().Debug("connection closed", zap.Uint64("clientId", info.Id))
				notifyClientOnline(info.UserId, info.Id)
			}, func(ctx context.Context, header *api_messages.MessageHeader, body io.Reader) {
				handleRpcMsg(info, ctx, header, body)
			}, setting.RpcMaxMessageSize)
			oldClient := GetClient(info.Id)
			clientMapLock.Lock()
			clients[info.Id] = client
			clientMapLock.Unlock()
			go func() {
				if oldClient != nil {
					oldClient.Close()
				}
				if !sendClientList(info, ctx) {
					CloseClient(info.UserId)
					return
				}
				notifyClientOnline(info.UserId, info.Id)
				err = client.HandleMessage()
				zap.L().Debug("HandleMessage exited", zap.Error(err), zap.Any("info", info))
			}()
		} else {
			zap.L().Error("failed to accept client connection", zap.Error(err), zap.String("clientId", info.Token), utils.UnknownError(ctx))
		}
	}
}

func sendAck(ctx context.Context, clientId uint64, messageId uint64, code api_messages.AckStatus) {
	err := sendMsgToClient(ctx, messageId, clientId, api_messages.MessageType_Ack, &api_messages.AckMessage{Code: code})
	if err != nil {
		zap.L().Debug("send ack failed", zap.Error(err), zap.Uint64("to", clientId), zap.Uint64("reply", messageId), zap.Any("code", code))
	}
}

func verifySender(usrId, from, to uint64) bool {
	ids := model.GetRelations(usrId)
	if ids == nil {
		return false
	}
	if _, ok := ids[from]; !ok {
		return false
	}
	if _, ok := ids[to]; !ok {
		return false
	}
	return true
}

func handleRpcMsg(info *model.TransportInfo, ctx context.Context, header *api_messages.MessageHeader, body io.Reader) {
	clientId := info.Id
	if !verifySender(info.UserId, clientId, header.ClientId) {
		sendAck(ctx, clientId, header.Id, api_messages.AckStatus_InvalidClientId)
		zap.L().Debug("invalid msg", zap.Uint64("form_client_id", clientId), zap.Uint64("to_client_id", clientId), zap.Uint64("msg_id", header.Id))
		return
	}
	msgClientId := header.ClientId
	client := GetClient(msgClientId)
	if client != nil {
		if ctx.Err() != nil {
			sendAck(ctx, clientId, header.Id, api_messages.AckStatus_MsgExpired)
			return
		}
		header.ClientId = clientId
		deadline, cancelFunc := context.WithDeadline(ctx, time.Unix(header.ExpiredAt, 0))
		defer cancelFunc()
		if err := client.SendMessage(deadline, header, body); err != nil {
			sendAck(ctx, clientId, header.Id, api_messages.AckStatus_SendFailed)
			zap.L().Debug("send message failed", zap.Error(err), zap.Uint64("to", msgClientId), zap.Uint64("from", clientId), zap.Uint64("msg_id", header.Id))
		} else {
			sendAck(ctx, clientId, header.Id, api_messages.AckStatus_Success)
			zap.L().Info("send message to client succeed", zap.Uint64("to", msgClientId), zap.Uint64("from", clientId), zap.Uint64("msg_id", header.Id))
		}
	} else {
		zap.L().Debug("client not found", zap.Uint64("to", msgClientId), zap.Uint64("form", clientId), zap.Uint64("msg_id", header.Id))
		sendAck(ctx, clientId, header.Id, api_messages.AckStatus_Offline)
	}
}

func GetClient(clientId uint64) (conn *socket.WebSocketConnection) {
	clientMapLock.RLock()
	defer clientMapLock.RUnlock()
	conn, _ = clients[clientId]
	return
}

func CloseClient(clientId uint64) {
	client := GetClient(clientId)
	if client != nil {
		client.Close()
	}
}

func DeleteTransport(userId uint64, token string) error {
	info, err := model.GetTransportInfoByTokenOrId(token, 0)
	if err != nil {
		return err
	}
	if info == nil || info.UserId != userId {
		return nil
	}
	err = model.DeleteTransportInfo(userId, token)
	if err != nil {
		return err
	}
	CloseClient(info.Id)
	return nil
}
